import pandas as pd
import numpy as np
import os
import argparse
from openpyxl import Workbook
from openpyxl.styles import PatternFill
from openpyxl.utils.dataframe import dataframe_to_rows


def extract_family_data(file_path='data.xlsx'):
    """
    提取家庭信息并生成新的Excel文件
    功能：提取游客姓名、性别、证件号码、联系方式、有无特殊要求及家属是否占床位等信息，并按家庭分组
    并将申请人标识改为显示申请人姓名
    
    参数:
        file_path: 输入的Excel文件路径，默认为'data.xlsx'
    """
    # 读取原Excel文件
    print(f"正在读取文件: {file_path}")
    
    try:
        # 使用dtype=str避免自动将数字转换为科学计数法
        df = pd.read_excel(file_path, dtype=str)
        print(f"文件读取成功！共 {len(df)} 行数据")
        
        # 创建结果列表，用于存储提取的个人信息
        result_data = []
        family_id = 1
        
        # 遍历每行数据（每个主申请人及其家属）
        for idx, row in df.iterrows():
            main_name = row['姓名'] if pd.notna(row['姓名']) else ''
            
            # 提取主申请人信息
            main_person = {
                '家庭ID': family_id,
                '申请人': main_name,  # 使用主申请人姓名
                '姓名': main_name,
                '性别': row['性别'] if pd.notna(row['性别']) else '',
                '证件号码': row['证件号码'] if pd.notna(row['证件号码']) else '',
                '联系方式': row['联系电话'] if pd.notna(row['联系电话']) else '',
                '有无特殊要求等等': row['有无特殊要求等等'] if pd.notna(row['有无特殊要求等等']) else '',  # 主申请人保留特殊要求
                '是否占床位': ''  # 主申请人不显示床位信息
            }
            result_data.append(main_person)
            
            # 提取家属信息（最多4位家属）
            for i in range(1, 5):
                family_name_col = f'家属{i}姓名'
                family_id_col = f'家属{i}身份证号码'
                
                # 检查家属姓名是否存在
                if family_name_col in row and pd.notna(row[family_name_col]) and row[family_name_col]:
                        family_gender_col = f'家属{i}性别'
                        family_bed_col = f'家属{i}是否占床位'
                        family_person = {
                            '家庭ID': family_id,
                            '申请人': main_name,  # 显示对应的主申请人姓名
                            '姓名': row[family_name_col],
                            '性别': row[family_gender_col] if pd.notna(row[family_gender_col]) else '',
                            '证件号码': row[family_id_col] if pd.notna(row[family_id_col]) else '',
                            '联系方式': row['联系电话'] if pd.notna(row['联系电话']) else '',  # 使用主申请人的联系方式
                            '有无特殊要求等等': '',  # 家属不显示特殊要求，避免重复
                            '是否占床位': row[family_bed_col] if family_bed_col in row and pd.notna(row[family_bed_col]) else ''
                        }
                        result_data.append(family_person)
            
            # 家庭ID递增
            family_id += 1
        
        # 创建结果DataFrame
        result_df = pd.DataFrame(result_data)
        
        # 按家庭ID排序，确保同一家人在一起
        # 不再需要按申请人标识排序，因为申请人信息已经在处理时放在前面了
        result_df = result_df.sort_values(by=['家庭ID'], ascending=[True])
        
        # 统计信息
        total_families = result_df['家庭ID'].nunique()
        total_persons = len(result_df)
        
        print(f"\n提取完成！")
        print(f"总家庭数: {total_families}")
        print(f"总人数: {total_persons}")
        
        # 定义用于标识不同家庭的颜色列表
        # 使用柔和的颜色，避免过于刺眼
        family_colors = [
            "EEEEEE",  # 浅灰色
            "DDEBF7",  # 浅蓝色
            "D5E8D4",  # 浅绿色
            "FFE6CC",  # 浅橙色
            "FFF2CC",  # 浅黄色
            "E1D5E7",  # 浅紫色
            "F8CECC",  # 浅红色
            "DAE8FC",  # 浅蓝色2
            "D5E8D4",  # 浅绿色2
            "FFF2CC"   # 浅黄色2
        ]
        
        # 生成新的Excel文件，使用新名称避免权限问题
        output_file = '家庭信息汇总_特殊要求合并.xlsx'
        
        # 创建一个新的工作簿
        wb = Workbook()
        ws = wb.active
        ws.title = "家庭信息汇总"
        
        # 将DataFrame的数据添加到工作表中
        for r_idx, row in enumerate(dataframe_to_rows(result_df, index=False, header=True), 1):
            for c_idx, value in enumerate(row, 1):
                ws.cell(row=r_idx, column=c_idx, value=value)
        
        # 为表头添加样式
        header_fill = PatternFill(start_color="4F81BD", end_color="4F81BD", fill_type="solid")
        for cell in ws[1]:
            cell.fill = header_fill
            cell.font = cell.font.copy(bold=True)
            cell.font = cell.font.copy(color="FFFFFF")
        
        # 为每个家庭分配颜色
        current_family_id = None
        color_index = 0
        
        # 从第2行开始（跳过表头）
        for row_idx in range(2, ws.max_row + 1):
            family_id = ws.cell(row=row_idx, column=1).value
            
            # 当家庭ID变化时，切换颜色
            if family_id != current_family_id:
                current_family_id = family_id
                # 使用循环的方式选择颜色，确保有足够的颜色
                color_index = (color_index + 1) % len(family_colors)
            
            # 为整行应用颜色
            fill = PatternFill(start_color=family_colors[color_index], end_color=family_colors[color_index], fill_type="solid")
            for cell in ws[row_idx]:
                cell.fill = fill
        
        # 合并同家庭的特殊要求单元格（如果有多个成员）
        # 先找出每家庭的起始和结束行
        family_rows = {}
        for row_idx in range(2, ws.max_row + 1):
            family_id = ws.cell(row=row_idx, column=1).value
            if family_id not in family_rows:
                family_rows[family_id] = {'start': row_idx, 'end': row_idx}
            else:
                family_rows[family_id]['end'] = row_idx
        
        # 找出特殊要求列的索引
        special_req_col_idx = result_df.columns.get_loc('有无特殊要求等等') + 1  # +1 因为Excel从1开始
        
        # 合并每个家庭的特殊要求单元格
        for family_id, rows in family_rows.items():
            # 只合并有多个成员的家庭，并且主申请人有特殊要求
            if rows['end'] > rows['start'] and ws.cell(row=rows['start'], column=special_req_col_idx).value:
                ws.merge_cells(start_row=rows['start'], start_column=special_req_col_idx,
                              end_row=rows['end'], end_column=special_req_col_idx)
        
        # 调整列宽以适应内容
        for column in ws.columns:
            max_length = 0
            column_letter = column[0].column_letter
            for cell in column:
                try:
                    if cell.value:
                        max_length = max(max_length, len(str(cell.value)))
                except:
                    pass
            adjusted_width = min(max_length + 2, 50)
            ws.column_dimensions[column_letter].width = adjusted_width
        
        # 保存工作簿
        wb.save(output_file)
        print(f"\n新Excel文件已生成: {os.path.abspath(output_file)}")
        print("不同家庭已用不同颜色标识，方便区分")
        print("特殊要求已按家庭合并在主申请人单元格中")
        print("已成功提取性别信息")
        print("已成功提取家属是否占床位信息")
        print("\n文件内容预览:")
        print(result_df.head(10))  # 显示前10行数据作为预览
        
    except Exception as e:
        print(f"处理文件时出错: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='提取家庭信息并生成新的Excel文件')
    # 添加文件路径参数，设置默认值为'data.xlsx'
    parser.add_argument('--file', '-f', default='data.xlsx', help='输入的Excel文件路径，默认为data.xlsx')
    # 解析命令行参数
    args = parser.parse_args()
    # 调用函数并传入文件路径参数
    extract_family_data(args.file)